from typing import Any, Iterable, Iterator, TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, SkipExhausted, pause, resume, skip
from zkl_aiutils_datasets.wrapping import wrap_dataset
from .sliced import SlicedDataset

AnyItem = TypeVar('AnyItem')


class BatchedIterator(DatasetIterator[tuple[AnyItem, ...]]):
    def __init__(self,
        iterator: Iterator[AnyItem],
        batch_items_n: int,
        keep_remainder: bool = False,
    ):
        self._iterator = iterator
        self._batch_items_n = batch_items_n
        self._keep_remainder = keep_remainder

    @property
    def iterator(self) -> Iterator[AnyItem]:
        return self._iterator

    @property
    def batch_items_n(self) -> int:
        return self._batch_items_n

    @property
    def keep_remainder(self) -> bool:
        return self._keep_remainder

    def __next__(self) -> tuple[AnyItem, ...]:
        batch_items = []
        for _ in range(self._batch_items_n):
            try:
                item = next(self._iterator)
                batch_items.append(item)
            except StopIteration:
                if self._keep_remainder and len(batch_items) > 0:
                    break
                raise
        return tuple(batch_items)

    def __skip__(self, batches_n: int):
        try:
            skip(self._iterator, batches_n * self._batch_items_n)
        except SkipExhausted as err:
            skipped_items_n = err.skipped_samples_n
            skipped_batches_n = skipped_items_n // self._batch_items_n
            if self._keep_remainder and skipped_items_n % self._batch_items_n:
                skipped_batches_n += 1
            raise SkipExhausted(skipped_batches_n)

    def __pause__(self) -> Any:
        return pause(self._iterator)


class BatchedDataset(Dataset[tuple[AnyItem, ...]]):
    def __init__(self,
        iterable: Iterable[AnyItem],
        batch_items_n: int,
        keep_remainder: bool = False,
    ):
        self._iterable = iterable
        self._dataset = wrap_dataset(iterable)
        self._batch_items_n = batch_items_n
        self._keep_remainder = keep_remainder

    @property
    def iterable(self) -> Iterable[AnyItem]:
        return self._iterable

    @property
    def batch_items_n(self) -> int:
        return self._batch_items_n

    @property
    def keep_remainder(self) -> bool:
        return self._keep_remainder

    def __len__(self) -> int:
        items_n = len(self._dataset)
        batches_n = items_n // self._batch_items_n
        if self._keep_remainder and items_n % self._batch_items_n:
            batches_n += 1
        return batches_n

    def __getitem__(self, batch_i: int) -> tuple[AnyItem, ...]:
        if batch_i < 0:
            raise IndexError(f'Expected 0<=batch_i, got {batch_i}.')

        try:
            items_n = len(self._dataset)
        except TypeError:
            items_n = None

        if items_n is not None:
            batches_n = items_n // self._batch_items_n
            if self._keep_remainder and items_n % self._batch_items_n:
                batches_n += 1
            if batch_i >= batches_n:
                raise IndexError(f'Expected batch_i<{batches_n}, got {batch_i}.')

        head = batch_i * self._batch_items_n
        tail = head + self._batch_items_n
        tail = tail if items_n is None else min(tail, items_n)
        return tuple(SlicedDataset(self._dataset, slice(head, tail)))

    def __iter__(self) -> BatchedIterator[AnyItem]:
        iterator = iter(self._iterable)
        return BatchedIterator(iterator, self._batch_items_n, self._keep_remainder)

    def __resume__(self, state: Any | None = None) -> BatchedIterator[AnyItem]:
        iterator = resume(self._iterable, state)
        return BatchedIterator(iterator, self._batch_items_n, self._keep_remainder)
